import pandas as pd
from snakemake.utils import min_version
min_version("6.0")

# Read the CSV file
samples_df = pd.read_csv('input.csv')  # Assuming columns are 'sample' and 'vireo_n'
SAMPLES = samples_df['sample'].tolist()

rule all:
    input:
        expand("{sample}/vireo", sample=SAMPLES)

rule index_bam:
    input:
        "{sample}/Aligned.sortedByCoord.out.bam"
    output:
        "{sample}/Aligned.sortedByCoord.out.bam.bai"
    threads: 16
    shell:
        """
        samtools index -@16 {input}
        """

rule cellsnp_lite:
    input:
        bam="{sample}/Aligned.sortedByCoord.out.bam",
        index="{sample}/Aligned.sortedByCoord.out.bam.bai"
    output:
        directory("{sample}/cellsnp")
    threads: 10
    shell:
        """
        cellsnp-lite -R ~/append-ssd/work/vcf/popscle_1kgp_0.01star.vcf.gz -p {threads} --minMAF 0.01 --minCOUNT 100 -s {input.bam} -b {wildcards.sample}/Solo.out/GeneFull_Ex50pAS/filtered_full/barcodes.tsv -O {output} --gzip
        """

rule vireosnp:
    input:
        "{sample}/cellsnp"
    output:
        directory("{sample}/vireo")
    threads: 8
    params:
        vireo_n=lambda wildcards: samples_df.loc[samples_df['sample'] == wildcards.sample, 'vireo_n'].values[0]
    shell:
        "vireo -p {threads} -c {input} -N {params.vireo_n} -o {output}"
